#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
NSGA-II算法实现模块
用于多目标优化问题求解
"""

import random
import time
from typing import Tuple

import numpy as np
import pandas as pd


class NSGA2:
    """NSGA-II算法实现"""

    def __init__(self, flag: int, generation: int, popsize: int, fjsp, optimizer, work: np.ndarray, order_num: int):
        """
        初始化NSGA-II算法
        
        Args:
            flag: 调度模式标识，0为普通调度，1为重调度
            generation: 迭代次数
            popsize: 种群规模
            fjsp: FJSP对象
            optimizer: 多目标优化对象
            work: 工作数据
            order_num: 订单数量
        """
        self.flag = flag
        self.generation = generation
        self.popsize = popsize
        self.fjsp = fjsp
        self.optimizer = optimizer
        self.work = work
        self.order_num = order_num

    def machine_crossover_mutation(self, wmt1: pd.DataFrame, wmt2: pd.DataFrame, current_gen: int) -> Tuple:
        """
        机器均匀交叉与变异
        
        Args:
            wmt1, wmt2: 待交叉的两个个体
            current_gen: 当前迭代次数
            
        Returns:
            交叉变异后的两个个体
        """
        # 初始化
        ccount = np.zeros((1, self.order_num), dtype=np.int32)
        var_prob = 1 - (current_gen / self.generation)

        # 按工件整理第二个个体
        wmt2_by_job = []
        for job in range(self.order_num):
            wmt2_by_job.append(wmt2[wmt2['W2'] == job])

        # 确定变异位点
        var_indices = []
        if var_prob > 0.5:
            var_indices = random.sample(range(wmt1.shape[0]), int(wmt1.shape[0] * 0.5 * var_prob))

        # 执行交叉和变异
        for i in wmt1.index:
            # 均匀交叉
            rand_num = np.random.randint(0, 2, 1)[0]
            w = int(wmt1.loc[i, 'W1'])

            if rand_num == 1:
                # 交换两个个体中相同位置的机器分配
                m = wmt1.loc[i, 'M1']
                t = wmt1.loc[i, 'T1']
                cross_index = wmt2_by_job[w].iloc[[ccount[0][w]]].index[0]
                wmt1.loc[i, 'M1'] = wmt2.loc[cross_index, 'M2']
                wmt1.loc[i, 'T1'] = wmt2.loc[cross_index, 'T2']
                wmt2.loc[cross_index, 'M2'] = m
                wmt2.loc[cross_index, 'T2'] = t

            # 变异操作
            if i in var_indices:
                if np.random.rand() > 0.5:
                    # 选取最小加工时间机器
                    min_machine = min(self.fjsp.info[w][ccount[0][w]].items(), key=lambda x: x[1])[0]
                    wmt1.loc[i, 'M1'] = min_machine
                    wmt1.loc[i, 'T1'] = self.fjsp.info[w][ccount[0, w]][min_machine]
                else:
                    # 随机挑选机器
                    n_machines = list(self.fjsp.info[w][ccount[0, w]].keys())
                    n_times = list(self.fjsp.info[w][ccount[0, w]].values())
                    index = np.random.randint(0, len(n_machines), 1)
                    wmt1.loc[i, 'M1'] = n_machines[index[0]]
                    wmt1.loc[i, 'T1'] = n_times[index[0]]

            ccount[0, w] = ccount[0][w] + 1

        # 转换为数组格式并返回
        W1 = np.array(wmt1['W1']).reshape(1, -1)
        m1 = np.array(wmt1['M1']).reshape(1, -1)
        t1 = np.array(wmt1['T1']).reshape(1, -1)
        W2 = np.array(wmt2['W2']).reshape(1, -1)
        m2 = np.array(wmt2['M2']).reshape(1, -1)
        t2 = np.array(wmt2['T2']).reshape(1, -1)

        return W1, m1, t1, W2, m2, t2

    def job_crossover(self, wmt1: pd.DataFrame, wmt2: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        工件位置交叉
        
        Args:
            wmt1, wmt2: 待交叉的两个个体
            
        Returns:
            交叉后的两个个体
        """
        # 生成随机顺序
        num = list(range(self.order_num))
        np.random.shuffle(num)
        index = np.random.randint(0, len(num), 1)[0]

        # 划分固定和变动的工件集合
        jpb_set1 = num[:index + 1]  # 固定不变的工件
        jpb_set2 = num[index + 1:]  # 按顺序读取的工件

        # 根据工件集合划分个体
        p1_in_s1 = wmt1[wmt1["W1"].isin(jpb_set1)]
        p2_in_s2 = wmt2[wmt2["W2"].isin(jpb_set2)]
        p2_in_s1 = wmt2[wmt2["W2"].isin(jpb_set1)]
        p1_in_s2 = wmt1[wmt1["W1"].isin(jpb_set2)]

        # 构建子代
        C1, C2 = [], []

        for i in range(wmt1.shape[0]):
            # 构造第一个子代
            if i in p1_in_s1.index:
                if p1_in_s1.shape[0] > 0:
                    C1.append(np.array(p1_in_s1.loc[i, :]))
            else:
                if p2_in_s2.shape[0] > 0:
                    C1.append(np.array(p2_in_s2.iloc[0, :]))
                    p2_in_s2 = p2_in_s2.iloc[1:, :]

            # 构造第二个子代
            if i in p2_in_s1.index:
                if p2_in_s1.shape[0] > 0:
                    C2.append(np.array(p2_in_s1.loc[i, :]))
            else:
                if p1_in_s2.shape[0] > 0:
                    C2.append(np.array(p1_in_s2.iloc[0, :]))
                    p1_in_s2 = p1_in_s2.iloc[1:, :]

        # 转换为DataFrame格式
        C1 = pd.DataFrame(C1)
        C1.columns = ['W1', 'M1', 'T1']
        C2 = pd.DataFrame(C2)
        C2.columns = ['W2', 'M2', 'T2']

        return C1, C2

    def crossover_and_mutation(self, wmt1: pd.DataFrame, wmt2: pd.DataFrame, current_gen: int) -> Tuple:
        """
        工件位置交叉和机器交叉的组合操作
        
        Args:
            wmt1, wmt2: 待交叉的两个个体
            current_gen: 当前迭代次数
            
        Returns:
            交叉变异后的两个个体
        """
        # 执行工序交叉
        start_time = time.time()
        C1, C2 = self.job_crossover(wmt1, wmt2)
        end_time = time.time()
        # print('工序交叉用时：', end_time - start_time)

        # 执行机器交叉
        start_time = time.time()
        W1, m1, t1, W2, m2, t2 = self.machine_crossover_mutation(C1, C2, current_gen)
        end_time = time.time()
        # print('机器交叉用时：', end_time - start_time)

        return W1, m1, t1, W2, m2, t2

    def run2(self) -> Tuple:
        """
        NSGA-II主循环
        用于重调度：
        优化目标：完工时间、负载均衡、重调度偏移率
        
        Returns:
            best_job, best_machine, best_time: 最优解
            fit_every: 迭代过程中的目标函数值
        """
        # 初始化
        answer = []
        fit_every = [[], [], [], []]

        # 初始化种群
        work_job1 = np.zeros((self.popsize, len(self.work)))
        work_M1 = np.zeros((self.popsize, len(self.work)))
        work_T1 = np.zeros((self.popsize, len(self.work)))

        work_job = np.zeros((self.popsize, len(self.work)))
        work_M = np.zeros((self.popsize, len(self.work)))
        work_T = np.zeros((self.popsize, len(self.work)))

        # 主循环
        for gen in range(self.generation):
            # 第一代：生成初始种群
            if gen < 1:
                for i in range(self.popsize):
                    # 生成个体
                    job, machine, machine_time = self.fjsp.encoding()

                    # 计算适应度
                    C_finish, LB, shift_rate, _, _, _, _, _ = self.fjsp.decoding2(job, machine, machine_time)

                    answer.append([C_finish, LB, shift_rate])
                    work_job[i], work_M[i], work_T[i] = job[0], machine[0], machine_time[0]

                # 计算分层和拥挤度
                # 调用三个目标的计算拥挤度
                _, crowder = self.optimizer.calculate_crowding_distance3(answer)

            # 按拥挤度排序并截取
            index_sort = crowder
            work_job = work_job[index_sort][0:self.popsize]
            work_M = work_M[index_sort][0:self.popsize]
            work_T = work_T[index_sort][0:self.popsize]
            answer = np.array(answer)[index_sort][0:self.popsize].tolist()

            # 生成子代
            answer1 = []

            for i in range(0, self.popsize, 2):
                # 选择父代
                W1, M1, T1 = work_job[i:i + 1], work_M[i:i + 1], work_T[i:i + 1]
                W2, M2, T2 = work_job[i + 1:i + 2], work_M[i + 1:i + 2], work_T[i + 1:i + 2]

                # 转换为DataFrame格式
                WMT1 = pd.DataFrame([W1[0], M1[0], T1[0]]).T
                WMT1.columns = ['W1', 'M1', 'T1']
                WMT2 = pd.DataFrame([W2[0], M2[0], T2[0]]).T
                WMT2.columns = ['W2', 'M2', 'T2']

                # 执行交叉变异
                C1, m1, t1, C2, m2, t2 = self.crossover_and_mutation(WMT1, WMT2, gen)

                # 计算子代适应度
                C_finish, LB, shift_rate, _, _, _, _, _ = self.fjsp.decoding2(C1, m1, t1)
                answer1.append([C_finish, LB, shift_rate])
                C_finish, LB, E_all, _, _, _, _, _ = self.fjsp.decoding2(C2, m2, t2)
                answer1.append([C_finish, LB, shift_rate])

                # 保存子代
                work_job1[i], work_M1[i], work_T1[i] = C1[0], m1[0], t1[0]
                work_job1[i + 1], work_M1[i + 1], work_T1[i + 1] = C2[0], m2[0], t2[0]

            # 合并父代和子代
            work_job = np.vstack((work_job, work_job1))
            work_M = np.vstack((work_M, work_M1))
            work_T = np.vstack((work_T, work_T1))
            answer = answer + answer1

            # 精英选择
            front, crowder = self.optimizer.calculate_crowding_distance3(answer)

            # 记录每代的Pareto前沿
            signal = front[0]
            pareto = np.array(answer)[signal]
            x = [pareto[i][0] for i in range(len(pareto))]
            y = [pareto[i][1] for i in range(len(pareto))]
            z = [pareto[i][2] for i in range(len(pareto))]

            fit_every[3].append(gen)
            fit_every[0].append([min(x), sum(x) / len(x), max(x)])
            fit_every[1].append([min(y), sum(y) / len(y), max(y)])
            fit_every[2].append([min(z), sum(z) / len(z), max(z)])

            # 选择最优解：距离远点最近的解
            min_pareto = np.argmin((pareto * pareto).sum(axis=1))
            pareto = pareto[min_pareto]
            best_index = [signal[min_pareto]]
            best_job, best_machine, best_time = work_job[best_index], work_M[best_index], work_T[best_index]

            print('算法迭代到了第%.0f次' % (gen + 1))

        return best_job, best_machine, best_time, fit_every

    def run1(self) -> Tuple:
        """
        NSGA-II主循环
        用于普通调度
        优化目标：完工时间、均衡负载

        Returns:
            best_job, best_machine, best_time: 最优解
            fit_every: 迭代过程中的目标函数值
        """
        # 初始化
        answer = []
        fit_every = [[], [], [], []]

        # 初始化种群
        work_job1 = np.zeros((self.popsize, len(self.work)))
        work_M1 = np.zeros((self.popsize, len(self.work)))
        work_T1 = np.zeros((self.popsize, len(self.work)))

        work_job = np.zeros((self.popsize, len(self.work)))
        work_M = np.zeros((self.popsize, len(self.work)))
        work_T = np.zeros((self.popsize, len(self.work)))

        start_time = time.time()

        # 主循环
        for gen in range(self.generation):
            # 第一代：生成初始种群
            if gen < 1:
                for i in range(self.popsize):
                    # 生成个体
                    job, machine, machine_time = self.fjsp.encoding()
                    C_finish, LB, _, _, _, _, _ = self.fjsp.decoding1(job, machine, machine_time)

                    # # 计算适应度
                    # if self.flag == 0:
                    #     C_finish, LB, E_all, _, _, _, _, _ = self.fjsp.decoding1(job, machine, machine_time)
                    # else:
                    #     C_finish, LB, E_all, _, _, _, _, _ = self.fjsp.decoding2(job, machine, machine_time)

                    # answer.append([C_finish, LB, E_all])
                    answer.append([C_finish, LB])

                    work_job[i], work_M[i], work_T[i] = job[0], machine[0], machine_time[0]

                end_time = time.time()
                # print('编码用时：', end_time - start_time)

                # 计算分层和拥挤度
                _, crowder = self.optimizer.calculate_crowding_distance(answer)

            # 按拥挤度排序并截取
            index_sort = crowder
            work_job = work_job[index_sort][0:self.popsize]
            work_M = work_M[index_sort][0:self.popsize]
            work_T = work_T[index_sort][0:self.popsize]
            answer = np.array(answer)[index_sort][0:self.popsize].tolist()

            # 生成子代
            answer1 = []
            start_time = time.time()

            for i in range(0, self.popsize, 2):
                # 选择父代
                W1, M1, T1 = work_job[i:i + 1], work_M[i:i + 1], work_T[i:i + 1]
                W2, M2, T2 = work_job[i + 1:i + 2], work_M[i + 1:i + 2], work_T[i + 1:i + 2]

                # 转换为DataFrame格式
                WMT1 = pd.DataFrame([W1[0], M1[0], T1[0]]).T
                WMT1.columns = ['W1', 'M1', 'T1']
                WMT2 = pd.DataFrame([W2[0], M2[0], T2[0]]).T
                WMT2.columns = ['W2', 'M2', 'T2']

                # 执行交叉变异
                C1, m1, t1, C2, m2, t2 = self.crossover_and_mutation(WMT1, WMT2, gen)
                C_finish, LB, _, _, _, _, _ = self.fjsp.decoding1(C1, m1, t1)
                answer1.append([C_finish, LB])
                C_finish, LB, _, _, _, _, _ = self.fjsp.decoding1(C2, m2, t2)
                answer1.append([C_finish, LB])

                # # 计算子代适应度
                # if self.flag == 0:
                #     C_finish, LB, E_all, _, _, _, _, _ = self.fjsp.decoding1(C1, m1, t1)
                #     answer1.append([C_finish, LB, E_all])
                #     C_finish, LB, E_all, _, _, _, _, _ = self.fjsp.decoding1(C2, m2, t2)
                #     answer1.append([C_finish, LB, E_all])
                # else:
                #     C_finish, LB, E_all, _, _, _, _, _ = self.fjsp.decoding2(C1, m1, t1)
                #     answer1.append([C_finish, LB, E_all])
                #     C_finish, LB, E_all, _, _, _, _, _ = self.fjsp.decoding2(C2, m2, t2)
                #     answer1.append([C_finish, LB, E_all])

                # 保存子代
                work_job1[i], work_M1[i], work_T1[i] = C1[0], m1[0], t1[0]
                work_job1[i + 1], work_M1[i + 1], work_T1[i + 1] = C2[0], m2[0], t2[0]

            end_time = time.time()
            # print('交叉变异：', end_time - start_time)

            # 合并父代和子代
            work_job = np.vstack((work_job, work_job1))
            work_M = np.vstack((work_M, work_M1))
            work_T = np.vstack((work_T, work_T1))
            answer = answer + answer1

            # 精英选择
            start_time = time.time()
            front, crowder = self.optimizer.calculate_crowding_distance(answer)
            end_time = time.time()
            # print('选择：', end_time - start_time)

            # 记录每代的Pareto前沿
            signal = front[0]
            pareto = np.array(answer)[signal]
            x = [pareto[i][0] for i in range(len(pareto))]
            y = [pareto[i][1] for i in range(len(pareto))]
            # z = [pareto[i][2] for i in range(len(pareto))]

            fit_every[3].append(gen)
            fit_every[0].append([min(x), sum(x) / len(x), max(x)])
            fit_every[1].append([min(y), sum(y) / len(y), max(y)])
            # fit_every[2].append([min(z), sum(z) / len(z), max(z)])

            # 选择最优解：距离远点最近的解
            min_pareto = np.argmin((pareto * pareto).sum(axis=1))
            pareto = pareto[min_pareto]
            best_index = [signal[min_pareto]]
            best_job, best_machine, best_time = work_job[best_index], work_M[best_index], work_T[best_index]

            print('算法迭代到了第%.0f次' % (gen + 1))

        return best_job, best_machine, best_time, fit_every
